import os
import re
import zipfile
from osgeo import gdal

os.environ['PROJ_LIB'] = r"/home/冯洋洋/PythonEnvironment/RasterShapeEnv/share/proj"

def update_band_array(l1c_dir):
    """"""
    for root, dir_names, file_names in os.walk(l1c_dir):
        for file_name in file_names:
            if re.match("T[\w]{5}_[\d]{8}T[\d]{6}_B[\w]{2}\.jp2$", file_name):
                jp2_path = os.path.join(root, file_name)
                new_path =  jp2_path.replace('.jp2', '-old.jp2')
                os.rename(jp2_path, new_path)
                attribute = get_raster_attribute(new_path)   
                arrays = [get_raster_band_array(new_path) * 0]
                get_raster_file(jp2_path, arrays, attribute)
                os.remove(new_path)

               


    # data to raster file
def get_raster_file(raster_path, arrays, attribute):
	""""""
	driver = gdal.GetDriverByName("GTiff")
	dataset = driver.Create(
		raster_path,
		xsize=attribute["RasterXSize"],
		ysize=attribute["RasterYSize"],
		bands=attribute["RasterCount"],
		eType=attribute["DataType"]
	)
	dataset.SetSpatialRef(attribute["SpatialRef"])
	dataset.SetGeoTransform(attribute["GeoTransform"])
	nodata = 0 if attribute["NoData"] is None else attribute["NoData"]
	for index in range(1, attribute["RasterCount"] + 1):
		band = dataset.GetRasterBand(index)
		band.WriteArray(arrays[index-1])
		band.SetNoDataValue(nodata)
	return


# raster to attribute
def get_raster_attribute(raster_path):
	""""""
	dataset = gdal.Open(raster_path)
	attributes = {
		"RasterXSize": dataset.__getattr__("RasterXSize"),
		"RasterYSize": dataset.__getattr__("RasterYSize"),
		"RasterCount": dataset.__getattr__("RasterCount"),
		"SpatialRef": dataset.GetSpatialRef(),
		"GeoTransform": dataset.GetGeoTransform(),
		"DataType": dataset.GetRasterBand(1).DataType,
		"NoData": dataset.GetRasterBand(1).GetNoDataValue(),
	}
	return attributes

# raster to array
def get_raster_band_array(raster_path):
	""""""
	dataset = gdal.Open(raster_path)
	xsize = dataset.RasterXSize
	ysize = dataset.RasterYSize
	array = dataset.ReadAsArray(
		xoff=0,
		yoff=0,
		xsize=xsize,
		ysize=ysize,
		interleave='band')
	return array


def main(png_dir, zip_dir, unzip_dir):
    for tile_dir_name in os.listdir(png_dir):
        tile_dir = os.path.join(png_dir, tile_dir_name)
        update_dir = os.path.join(tile_dir, "update")

        for png_name in os.listdir(tile_dir):
            if png_name != "update":
                zip_name = png_name.replace(".png", ".zip")
                zip_path = os.path.join(zip_dir, zip_name)
                with zipfile.ZipFile(zip_path, "r") as zip_file:
                    for fileInfo in zip_file.infolist():
                        zip_file.extract(fileInfo.filename, unzip_dir)

        for png_name in os.listdir(update_dir):
            zip_name = png_name.replace(".png", ".zip")
            zip_path = os.path.join(zip_dir, zip_name)
            with zipfile.ZipFile(zip_path, "r") as zip_file:
                for fileInfo in zip_file.infolist():
                    zip_file.extract(fileInfo.filename, unzip_dir)
                    safe_name = png_name.replace(".png", ".SAFE")
                    safe_dir = os.path.join(unzip_dir, safe_name)
                    update_band_array(safe_dir)
        
    return

if __name__ == "__main__":
    png_dir = r"/home/遥感技术部/遥感数据/冯洋洋/png"
    zip_dir = r"/home/遥感技术部/遥感数据/冯洋洋/test"
    unzip_dir = r"/home/遥感技术部/遥感数据/冯洋洋/unzip"
    main(png_dir, zip_dir, unzip_dir)